package com.tre.jdevtemplateboot.web.filter;

import com.tre.jdevtemplateboot.common.util.XSSWhiteListUtils;
import org.springframework.util.StringUtils;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    HttpServletRequest orgRequest = null;
    private Map<String, String[]> parameterMap = new HashMap<>();

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        orgRequest = request;
        parameterMap.putAll(request.getParameterMap());
    }


    @Override
    public String getHeader(String name) {
        String strHeader = super.getHeader(name);
        if (!StringUtils.hasLength(strHeader)) {
            return strHeader;
        }
        return XSSWhiteListUtils.clean(strHeader);
    }

    @Override
    public String getParameter(String name) {
        String strParameter = super.getParameter(name);
        if (!StringUtils.hasLength(strParameter)) {
            return strParameter;
        }
        return XSSWhiteListUtils.clean(strParameter);
    }


    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values == null) {
            return new String[0];
        }
        int length = values.length;
        String[] escapesValues = new String[length];
        for (int i = 0; i < length; i++) {
            escapesValues[i] = XSSWhiteListUtils.clean(values[i]);
        }
        return escapesValues;
    }

    @Override
    public Map<String, String[]> getParameterMap() {

        for (Iterator iterator = parameterMap.entrySet().iterator(); iterator.hasNext(); ) {
            Map.Entry entry = (Map.Entry) iterator.next();
            String[] values = (String[]) entry.getValue();
            for (int i = 0; i < values.length; i++) {
                if (values[i] instanceof String) {
                    values[i] = XSSWhiteListUtils.clean(values[i]);
                }
            }
        }
        return parameterMap;
    }

    @Override
    public ServletInputStream getInputStream() {
        String body = HttpGetBody.getBodyString(this.getOrgRequest());
        ServletInputStream inputStream = null;
        if (StringUtils.hasLength(body)) {
            body = XSSWhiteListUtils.clean(body);
            inputStream = new TranslateServletInputStream(body);
        }
        return inputStream;
    }

    //获取最原始的request
    public HttpServletRequest getOrgRequest() {
        return orgRequest;
    }

}
